package middleware

import (
	"bytes"
	"github.com/gin-gonic/gin"
	"io"
	"log"
	"strings"
	"time"
)

type responseWriter struct {
	gin.ResponseWriter
	b *bytes.Buffer
}

func (w responseWriter) Write(b []byte) (int, error) {
	w.b.Write(b)
	return w.ResponseWriter.Write(b)
}

// LogMiddleware 日志记录中间件
func LogMiddleware(ctx *gin.Context) {
	// requests 请求前
	startTime := time.Now()
	traceId := ctx.Request.Header.Get("traceId")
	reqType := ctx.Request.Header.Get("Content-Type")
	if reqType == "application/json" {
		url := ctx.Request.RequestURI
		data, _ := ctx.GetRawData()
		ctx.Request.Body = io.NopCloser(bytes.NewBuffer(data))
		log.Printf("请求消息：traceId: %v, path: %v, data: %v\n", traceId, url, string(data))
	}

	writer := responseWriter{
		ctx.Writer,
		bytes.NewBuffer([]byte{}),
	}
	ctx.Writer = writer
	ctx.Next()
	resType := ctx.Writer.Header().Get("Content-Type")
	if strings.Contains(resType, "application/json") {
		duration := time.Now().Sub(startTime)
		responseBody := writer.b.String()
		log.Printf("返回消息：traceId: %v, durtion: %v, response_body: %v\n", traceId, duration, responseBody)
	}
}
